
import os
import random
from random import shuffle

import cv2
import numpy as np
from PIL import Image
from torch.utils.data.dataset import Dataset


def rand(a=0, b=1):
    return np.random.rand()*(b-a) + a

def get_random_data(image, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, flip_signal=False):

    h, w = input_shape
    # ------------------------------------------#
    #   图像大小调整
    # ------------------------------------------#
    rand_jit1 = rand(1 - jitter, 1 + jitter)
    rand_jit2 = rand(1 - jitter, 1 + jitter)
    new_ar = w / h * rand_jit1 / rand_jit2  # 变化后图片的长宽比例

    scale = rand(0.75, 1.25)  # 这里就是图片大小在75%到125%之间
    if new_ar < 1:
        nh = int(scale * h)
        nw = int(nh * new_ar)
    else:
        nw = int(scale * w)
        nh = int(nw / new_ar)
    image = image.resize((nw, nh), Image.BICUBIC)  # 三次样条插值

    # ------------------------------------------#
    #   翻转图像
    # ------------------------------------------#
    flip = rand() < .5  # 按照概率翻转
    if flip and flip_signal:
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
    # ------------------------------------------#
    #   随机遮盖图像
    # ------------------------------------------#
    dx = int(rand(0, w - nw))
    dy = int(rand(0, h - nh))
    new_image = Image.new('RGB', (w, h), (255, 255, 255))
    new_image.paste(image, (dx, dy))
    image = new_image
    # ------------------------------------------#
    #   图像旋转
    # ------------------------------------------#
    rotate = rand() < .5
    if rotate:
        angle = np.random.randint(-5, 5)
        a, b = w / 2, h / 2
        M = cv2.getRotationMatrix2D((a, b), angle, 1)
        image = cv2.warpAffine(np.array(image), M, (w, h), borderValue=[255, 255, 255])
    # ------------------------------------------#
    #   色域扭曲 转到HSV色域调整
    # ------------------------------------------#
    hue = rand(-hue, hue)
    sat = rand(1, sat) if rand() < .5 else 1 / rand(1, sat)
    val = rand(1, val) if rand() < .5 else 1 / rand(1, val)
    x = cv2.cvtColor(np.array(image, np.float32) / 255, cv2.COLOR_RGB2HSV)
    x[..., 0] += hue * 360
    x[..., 0][x[..., 0] > 1] -= 1
    x[..., 0][x[..., 0] < 0] += 1
    x[..., 1] *= sat
    x[..., 2] *= val
    x[x[:, :, 0] > 360, 0] = 360
    x[:, :, 1:][x[:, :, 1:] > 1] = 1
    x[x < 0] = 0
    image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB) * 255

    return image_data

image = Image.open("./img/1.jpg")
input_shape = [448,448]
# image = np.asarray(image).astype(np.float64)
# image = np.transpose(image, [2, 0, 1])  # 变换通道
# image = image / 255  # 将值置0，1
get_random_data(image,input_shape)